#include "conv_1x1.h"

namespace mslite {

Conv1x1::Conv1x1(cl_context context, cl_device_id device, std::string clKernelName, std::vector<int> ishape, std::vector<int> wshape): Operation(context) {
    this->clFileName = clKernelName + ".cl";
    this->clKernelName = clKernelName;
    this->context = context;
    this->inputNames = {"input", "weight"};
    this->inputShapes = {ishape, wshape};                                   // ishape must be nhwc, {1, 36, 14400, 128};
    this->outputNames = {"output"};
    std::vector<int> oshape = {ishape[0], ishape[1], ishape[2], wshape[0]};   // oshape is also nhwc, {1, 36, 14400, 256};
    this->outputShapes = {oshape};

    readSrcCode(this->clFileName, this->clCode);
    compileKernel(this->clCode, this->clCode.size(), this->clKernelName, this->context, device, this->program, this->kernel);
    // prepareMem();
}

cl_int Conv1x1::prepareMem(void) {
    cl_int status;
    memObjs.clear();
    int inputSize = getTotalSize(inputShapes[0]);
    int weightSize = getTotalSize(inputShapes[1]);
    int outputSize = getTotalSize(outputShapes[0]);
    int biasSize = inputShapes[1][0];
    memObjs["input"] = clCreateBuffer(context, CL_MEM_READ_ONLY |  CL_MEM_ALLOC_HOST_PTR, inputSize*sizeof(float), NULL, &status);
    memObjs["weight"] = clCreateBuffer(context, CL_MEM_READ_ONLY |  CL_MEM_ALLOC_HOST_PTR, weightSize*sizeof(float), NULL, &status);
    memObjs["output"] = clCreateBuffer(context, CL_MEM_WRITE_ONLY | CL_MEM_ALLOC_HOST_PTR, outputSize*sizeof(float), NULL, &status);
    std::vector<float> bias(biasSize, 0);
    memObjs["bias"] = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR, biasSize*sizeof(float), bias.data(), &status);

    memTypes["input"] = Buffer;
    memTypes["weight"] = Buffer;
    memTypes["output"] = Buffer;
    memTypes["bias"] = Buffer;

    return status;
}

cl_int Conv1x1::setParams(void) {
    globalThreads.resize(3);
    localThreads.resize(3);
    if (clKernelName == "conv_1x1_tflite")
        globalThreads = {static_cast<size_t>(outputShapes[0][2]/2), static_cast<size_t>(outputShapes[0][1]), static_cast<size_t>(outputShapes[0][3]/8)};
    else 
        globalThreads = {static_cast<size_t>(outputShapes[0][2]/2), static_cast<size_t>(outputShapes[0][1]), static_cast<size_t>(outputShapes[0][3]/8)};
    // std::cout << globalThreads[0] << " " << globalThreads[1] << " " << globalThreads[2] << std::endl;
    localThreads = {8, 4, 1};
    return 0;
}

// cl_int Conv1x1::getOutut(cl_command_queue commandQueue, std::vector<std::vector<float>> &output) {
//     cl_int status;
//     int osize = getTotalSize(this->outputShapes[0]);
//     output.clear();
//     output.resize(1);
//     output[0].resize(osize);
//     status = mslite::clEnqueueReadBuffer(commandQueue, memObjs["output"], CL_TRUE, 0, sizeof(float)*osize, output[0].data(), 0, NULL, NULL);
//     return status;
// }

Conv1x1Naive::Conv1x1Naive(cl_context context, cl_device_id device, std::vector<int> ishape, std::vector<int> wshape): 
    Conv1x1(context, device, "conv_1x1_naive", ishape, wshape) {
}

cl_int Conv1x1Naive::bind(void) {
    cl_int status;

    cl_int4 input_sliced_shape[] = {inputShapes[0][0], inputShapes[0][1], inputShapes[0][2], inputShapes[0][3]/4};
    cl_int4 weight_sliced_shape[] = {inputShapes[1][0]/4, inputShapes[1][1], inputShapes[1][2], inputShapes[1][3]/4};

    status = clSetKernelArg(kernel, 0, sizeof(cl_mem), static_cast<void *>(&memObjs["output"]));
    status = clSetKernelArg(kernel, 1, sizeof(cl_mem), static_cast<void *>(&memObjs["input"]));
    status = clSetKernelArg(kernel, 2, sizeof(cl_mem), static_cast<void *>(&memObjs["weight"]));
    status = clSetKernelArg(kernel, 3, sizeof(cl_int4), static_cast<void *>(input_sliced_shape));
    status = clSetKernelArg(kernel, 4, sizeof(cl_int4), static_cast<void *>(weight_sliced_shape));
    return status;
}

Conv1x1Tflite::Conv1x1Tflite(cl_context context, cl_device_id device, std::vector<int> ishape, std::vector<int> wshape):
    Conv1x1(context, device, "conv_1x1_tflite", ishape, wshape) {
}

cl_int Conv1x1Tflite::bind(void) {
    cl_int status;

    cl_int4 a[] = {outputShapes[0][2], outputShapes[0][1], outputShapes[0][3]/4, outputShapes[0][1]};
    cl_int4 b[] = {inputShapes[0][3]/4, inputShapes[0][2], 0, 0};

    status = clSetKernelArg(kernel, 0, sizeof(cl_mem), static_cast<void *>(&memObjs["bias"]));
    status = clSetKernelArg(kernel, 1, sizeof(cl_mem), static_cast<void *>(&memObjs["output"]));
    status = clSetKernelArg(kernel, 2, sizeof(cl_mem), static_cast<void *>(&memObjs["input"]));
    status = clSetKernelArg(kernel, 3, sizeof(cl_mem), static_cast<void *>(&memObjs["weight"]));
    status = clSetKernelArg(kernel, 4, sizeof(cl_int4), static_cast<void *>(a));
    status = clSetKernelArg(kernel, 5, sizeof(cl_int4), static_cast<void *>(b));
    return status;
}

Conv1x1Image2d::Conv1x1Image2d(cl_context context, cl_device_id device, std::vector<int> ishape, std::vector<int> wshape):
    Conv1x1(context, device, "conv_1x1_image2d", ishape, wshape) {

}

cl_int Conv1x1Image2d::prepareMem(void) {
    cl_int status;
    memObjs.clear();

    int ih = inputShapes[0][1] * inputShapes[0][3] / 4, iw = inputShapes[0][2];
    int oh = outputShapes[0][1] * outputShapes[0][3] / 4, ow = outputShapes[0][2];
    memObjs["input"] = createImage2d(context, CL_FLOAT, iw, ih, NULL, &status);
    memObjs["output"] = createImage2d(context, CL_FLOAT, ow, oh, NULL, &status);

    int weightSize = getTotalSize(inputShapes[1]);
    memObjs["weight"] = clCreateBuffer(context, CL_MEM_READ_ONLY |  CL_MEM_ALLOC_HOST_PTR, weightSize*sizeof(float), NULL, &status);

    memTypes["input"] = Image2d;
    memTypes["weight"] = Buffer;
    memTypes["output"] = Image2d;

    return status;
}

cl_int Conv1x1Image2d::bind(void) {
    cl_int status;

    cl_int4 input_sliced_shape[] = {inputShapes[0][0], inputShapes[0][1], inputShapes[0][2], inputShapes[0][3]/4};
    cl_int4 weight_sliced_shape[] = {inputShapes[1][0]/4, inputShapes[1][1], inputShapes[1][2], inputShapes[1][3]/4};

    status = clSetKernelArg(kernel, 0, sizeof(cl_mem), static_cast<void *>(&memObjs["output"]));
    status = clSetKernelArg(kernel, 1, sizeof(cl_mem), static_cast<void *>(&memObjs["input"]));
    status = clSetKernelArg(kernel, 2, sizeof(cl_mem), static_cast<void *>(&memObjs["weight"]));
    status = clSetKernelArg(kernel, 3, sizeof(cl_int4), static_cast<void *>(input_sliced_shape));
    status = clSetKernelArg(kernel, 4, sizeof(cl_int4), static_cast<void *>(weight_sliced_shape));
    return status;
}

}